<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>gpu computation renderer</title>
  <style>
    body {
      margin: 0;
      padding: 0;
      overflow: hidden;
    }
  </style>
  <script id="drawVert" type="x-shader/x-fragment">
    precision mediump float;

    attribute float a_index;

    uniform sampler2D u_particles;
    uniform float u_particles_res;

    varying vec2 v_particle_pos;

    void main() {
      vec4 color = texture2D(u_particles, vec2(
      fract(a_index / u_particles_res),
      floor(a_index / u_particles_res) / u_particles_res));

      // decode current particle position from the pixel's RGBA value
      v_particle_pos = vec2(
      color.r / 255.0 + color.b,
      color.g / 255.0 + color.a);

      gl_PointSize = 1.0;
      vec2 pos = vec2(2.0 * v_particle_pos.x - 1.0, 1.0 - 2.0 * v_particle_pos.y);
      gl_Position = vec4(pos, 0, 1);
    }
  </script>
  <script id="drawFrag" type="x-shader/x-fragment">
    precision mediump float;

    uniform sampler2D u_wind;
    uniform vec2 u_wind_min;
    uniform vec2 u_wind_max;
    uniform sampler2D u_color_ramp;

    varying vec2 v_particle_pos;

    void main() {
      vec2 velocity = mix(u_wind_min, u_wind_max, texture2D(u_wind, v_particle_pos).rg);
      float speed_t = length(velocity) / length(u_wind_max);

      // color ramp is encoded in a 16x16 texture
      vec2 ramp_pos = vec2(
      fract(16.0 * speed_t),
      floor(16.0 * speed_t) / 16.0);

      gl_FragColor = texture2D(u_color_ramp, ramp_pos);
      // gl_FragColor = texture2D(u_wind, v_particle_pos);
    }
  </script>
  <script id="quadVert" type="x-shader/x-fragment">
    precision mediump float;

    varying vec2 v_tex_pos;

    void main() {
      v_tex_pos = uv;
      gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
    }
  </script>
  <script id="screenFrag" type="x-shader/x-fragment">
    precision mediump float;

    uniform sampler2D u_screen;
    uniform float u_opacity;

    varying vec2 v_tex_pos;

    void main() {
      vec4 color = texture2D(u_screen, 1.0 - v_tex_pos);
      // a hack to guarantee opacity fade out even with a value close to 1.0
      gl_FragColor = vec4(floor(255.0 * color * u_opacity) / 255.0);
    }
  </script>
  <script id="fragmentShaderVelocity" type="x-shader/x-fragment">
    precision highp float;

//    uniform sampler2D u_particles;
    uniform sampler2D u_wind;
    uniform vec2 u_wind_res;
    uniform vec2 u_wind_min;
    uniform vec2 u_wind_max;

    varying vec2 v_tex_pos;

    // wind speed lookup; use manual bilinear filtering based on 4 adjacent pixels for smooth interpolation
    vec2 lookup_wind(const vec2 uv) {
      // return texture2D(u_wind, uv).rg; // lower-res hardware filtering
      vec2 px = 1.0 / u_wind_res;
      vec2 vc = (floor(uv * u_wind_res)) * px;
      vec2 f = fract(uv * u_wind_res);
      vec2 tl = texture2D(u_wind, vc).rg;
      vec2 tr = texture2D(u_wind, vc + vec2(px.x, 0)).rg;
      vec2 bl = texture2D(u_wind, vc + vec2(0, px.y)).rg;
      vec2 br = texture2D(u_wind, vc + px).rg;
      return mix(mix(tl, tr, f.x), mix(bl, br, f.x), f.y);
    }

    void main() {
      vec4 color = texture2D(u_particles, v_tex_pos);
      vec2 pos = vec2(
      color.r / 255.0 + color.b,
      color.g / 255.0 + color.a); // decode particle position from pixel RGBA

      vec2 velocity = mix(u_wind_min, u_wind_max, lookup_wind(pos));
      velocity = velocity / length(u_wind_max);
      gl_FragColor = vec4(velocity, 0.0, 0.0);
    }
  </script>
  <script id="fragmentShaderPosition" type="x-shader/x-fragment">
    precision highp float;

//    uniform sampler2D u_particles;
//    uniform sampler2D u_velocity;
    uniform vec2 u_wind_max;
    uniform float u_rand_seed;
    uniform float u_speed_factor;
    uniform float u_drop_rate;
    uniform float u_drop_rate_bump;

    varying vec2 v_tex_pos;

    // pseudo-random generator
    const vec3 rand_constants = vec3(12.9898, 78.233, 4375.85453);
    float rand(const vec2 co) {
      float t = dot(rand_constants.xy, co);
      return fract(sin(t) * (rand_constants.z + t));
    }

    void main() {
      vec4 velocity = texture2D(u_velocity, v_tex_pos);

      float speed_t = length(velocity);
      velocity *= length(u_wind_max);

      vec4 color = texture2D(u_particles, v_tex_pos);
      vec2 pos = vec2(
      color.r / 255.0 + color.b,
      color.g / 255.0 + color.a); // decode particle position from pixel RGBA
      // take EPSG:4236 distortion into account for calculating where the particle moved
      float distortion = cos(radians(pos.y * 180.0 - 90.0));
      // https://discourse.threejs.org/t/r178-three-js-ping-pong-rendering-texture-issue/85839/8
      // add 0.000001 at end to avoid divide zero exception
      distortion = abs(distortion) - 0.000001;
      vec2 offset = vec2(velocity.x / distortion, -velocity.y) * 0.0001 * u_speed_factor;

      // update particle position, wrapping around the date line
      pos = fract(1.0 + pos + offset);

      // a random seed to use for the particle drop
      vec2 seed = (pos + v_tex_pos) * u_rand_seed;

      // drop rate is a chance a particle will restart at random position, to avoid degeneration
      float drop_rate = u_drop_rate + speed_t * u_drop_rate_bump;
      float drop = step(1.0 - drop_rate, rand(seed));

      vec2 random_pos = vec2(
      rand(seed + 1.3),
      rand(seed + 2.1));
      pos = mix(pos, random_pos, drop);

      // encode the new particle position back into RGBA
      gl_FragColor = vec4(
      fract(pos * 255.0),
      floor(pos * 255.0) / 255.0);
    }
  </script>
</head>
<body>
<script type="importmap">
  {
    "imports": {
      "three": "https://cdn.jsdelivr.net/npm/three@0.179.1/build/three.module.js",
      "three/addons/": "https://cdn.jsdelivr.net/npm/three@0.179.1/examples/jsm/"
    }
  }
</script>
<script type="module">
  import * as THREE from 'three'
  import {OrbitControls} from "three/addons/controls/OrbitControls.js"
  // import {GPUComputationRenderer} from "three/addons/misc/GPUComputationRenderer.js";
  // I changed a few code in GPUComputationRenderer.js
  import {GPUComputationRenderer} from "../lib/jsm/misc/GPUComputationRenderer.js";

  const drawVert = document.getElementById("drawVert").textContent;
  const drawFrag = document.getElementById("drawFrag").textContent;
  const quadVert = document.getElementById("quadVert").textContent;
  const screenFrag = document.getElementById("screenFrag").textContent;
  // webgl_framebuffer_texture webgl_rtt
  function getJSON(url, callback) {
    const xhr = new XMLHttpRequest();
    xhr.responseType = 'json';
    xhr.open('get', url, true);
    xhr.onload = function () {
      if (xhr.status >= 200 && xhr.status < 300) {
        callback(xhr.response);
      } else {
        throw new Error(xhr.statusText);
      }
    };
    xhr.send();
  }

  const defaultRampColors = {
    0.0: '#3288bd',
    0.1: '#66c2a5',
    0.2: '#abdda4',
    0.3: '#e6f598',
    0.4: '#fee08b',
    0.5: '#fdae61',
    0.6: '#f46d43',
    1.0: '#d53e4f'
  };

  function getColorRamp(colors) {
    const canvas = document.createElement('canvas');
    const ctx = canvas.getContext('2d');

    canvas.width = 256;
    canvas.height = 1;

    const gradient = ctx.createLinearGradient(0, 0, 256, 0);
    for (const stop in colors) {
      gradient.addColorStop(+stop, colors[stop]);
    }

    ctx.fillStyle = gradient;
    ctx.fillRect(0, 0, 256, 1);

    return new Uint8Array(ctx.getImageData(0, 0, 256, 1).data);
  }

  const renderer = new THREE.WebGLRenderer( { antialias: false } );
  renderer.setPixelRatio( window.devicePixelRatio );
  renderer.setSize( window.innerWidth, window.innerHeight );
  renderer.outputColorSpace = THREE.LinearSRGBColorSpace
  renderer.autoClear = true;
  renderer.setClearColor( 0x000000, 1 );
  document.body.appendChild( renderer.domElement );

  const scene = new THREE.Scene();
  const sceneBack = new THREE.Scene();
  const sceneFront = new THREE.Scene();

  const width = window.innerWidth;
  const height = window.innerHeight;
  const camera = new THREE.PerspectiveCamera( 45, width / height, 0.01, 200 );
  camera.position.set( 0, 0, 1.42 );
  camera.up.set(0, 0, 1);
  camera.lookAt( 0, 0, 0 );

  const controls = new OrbitControls( camera, renderer.domElement );
  // controls.maxPolarAngle = Math.PI * 0.5;
  controls.minDistance = 0.01;
  controls.maxDistance = 100;
  // controls.addEventListener( 'change', render );

  scene.add( new THREE.AxesHelper(2000) )

  const planeGeometry = new THREE.PlaneGeometry( 2, 1, 1 );
  const planeMaterial = new THREE.MeshBasicMaterial( {wireframe: false, color: 0xffffff, side: THREE.DoubleSide } );
  const planeMesh = new THREE.Mesh( planeGeometry, planeMaterial );
  scene.add(planeMesh);
  planeMesh.position.set( 0, 0, 0.001 );

  getJSON('./assets/ne_110m_coastline.geojson', function (data) {
    const positions = [];
    const geometry = new THREE.BufferGeometry();
    const material = new THREE.LineBasicMaterial( { color: '#fff' } );
    for (let i = 0; i < data.features.length; i++) {
      const line = data.features[i].geometry.coordinates;
      for (let j = 0; j < line.length; j++) {
        const x = (line[j][0] + 180) * 2 / 360;
        const y = (line[j][1] + 90) * 1 / 180;
        positions.push(x, y, 0)
        if (j && j < line.length - 1) {
          positions.push(x, y, 0)
        }
      }
    }
    geometry.setAttribute( 'position', new THREE.Float32BufferAttribute( positions, 3 ) );
    const contours = new THREE.LineSegments(geometry, material );
    contours.position.set( -1, -0.5, 0 );
    scene.add( contours );
  });

  const windFiles = {
    0: '2016112000',
    6: '2016112006',
    12: '2016112012',
    18: '2016112018',
    24: '2016112100',
    30: '2016112106',
    36: '2016112112',
    42: '2016112118',
    48: '2016112200'
  };

  async function updateWind(name) {
    return new Promise( ( resolve, reject ) => {
      getJSON('./assets/wind/' + windFiles[name] + '.json', function (windData) {
        const windImage = new Image();
        windData.image = windImage;
        windImage.src = './assets/wind/' + windFiles[name] + '.png';
        windImage.onload = function () {

          resolve(windData);
        };
      });
    })
  }
  const windFileKey = 0
  const windData = await updateWind(windFileKey)
  console.log(windData)

  const particleRes = 256
  const _numParticles = particleRes * particleRes;
  const textureColors = new THREE.DataTexture(getColorRamp(defaultRampColors), 16, 16, THREE.RGBAFormat, THREE.UnsignedByteType);
  textureColors.magFilter = THREE.LinearFilter;
  textureColors.minFilter = THREE.LinearFilter;
  textureColors.generateMipmaps = false;
  textureColors.needsUpdate = true;

  const windTexture = new THREE.Texture(windData.image);
  windTexture.magFilter = THREE.LinearFilter;
  windTexture.minFilter = THREE.LinearFilter;
  windTexture.wrapS = THREE.ClampToEdgeWrapping;
  windTexture.wrapT = THREE.ClampToEdgeWrapping;
  windTexture.format = THREE.RGBAFormat;
  windTexture.needsUpdate = true;
  const WindParam = {
    fadeOpacity: 0.996, // how fast the particle trails fade on each frame
    speedFactor: 0.25, // how fast the particles move
    dropRate: 0.003, // how often the particles move to a random place
    dropRateBump: 0.01, // drop rate increase relative to individual particle speed
  }

  let gpuCompute, velocityVariable, positionVariable, windVariable;
  let positionUniforms, velocityUniforms;

  function initComputeRenderer() {
    gpuCompute = new GPUComputationRenderer( particleRes, particleRes, renderer );
    const particleState = new Uint8Array(_numParticles * 4);
    for (let i = 0; i < particleState.length; i++) {
      particleState[i] = Math.floor(Math.random() * 256); // randomize the initial particle positions
    }
    const dtPosition = new THREE.DataTexture( particleState, particleRes, particleRes, THREE.RGBAFormat, THREE.UnsignedByteType );
    dtPosition.needsUpdate = true;
    const dtVelocity = new THREE.DataTexture( particleState, particleRes, particleRes, THREE.RGBAFormat, THREE.UnsignedByteType );
    dtVelocity.needsUpdate = true;
    velocityVariable = gpuCompute.addVariable( 'u_velocity', document.getElementById( 'fragmentShaderVelocity' ).textContent, dtVelocity );
    positionVariable = gpuCompute.addVariable( 'u_particles', document.getElementById( 'fragmentShaderPosition' ).textContent, dtPosition );

    gpuCompute.setVariableDependencies( velocityVariable, [ positionVariable, velocityVariable ] );
    gpuCompute.setVariableDependencies( positionVariable, [ positionVariable, velocityVariable ] );

    positionUniforms = positionVariable.material.uniforms;
    velocityUniforms = velocityVariable.material.uniforms;

    velocityUniforms[ 'u_wind_res' ] = { value: {x: windData.width, y: windData.height} };
    velocityUniforms[ 'u_wind' ] = { value: windTexture };
    velocityUniforms[ 'u_wind_min' ] = { value: {x: windData.uMin, y: windData.vMin} };
    velocityUniforms[ 'u_wind_max' ] = { value: {x: windData.uMax, y: windData.vMax} };

    positionUniforms[ 'u_wind_max' ] = { value: {x: windData.uMax, y: windData.vMax} };
    positionUniforms[ 'u_rand_seed' ] = { value: Math.random() };
    positionUniforms[ 'u_speed_factor' ] = { value: WindParam.speedFactor };
    positionUniforms[ 'u_drop_rate' ] = { value: WindParam.dropRate };
    positionUniforms[ 'u_drop_rate_bump' ] = { value: WindParam.dropRateBump };

    velocityVariable.wrapS = THREE.ClampToEdgeWrapping;
    velocityVariable.wrapT = THREE.ClampToEdgeWrapping;
    positionVariable.wrapS = THREE.ClampToEdgeWrapping;
    positionVariable.wrapT = THREE.ClampToEdgeWrapping;
    const error = gpuCompute.init();

    if ( error !== null ) {
      console.error( error );
    }
    // planeMaterial.map = dtPosition;
  }
  initComputeRenderer();

  const opt = {
    magFilter: THREE.NearestFilter,
    minFilter: THREE.NearestFilter,
    format: THREE.RGBAFormat,
    type: THREE.UnsignedByteType,
    depthBuffer: false
  }
  const screenTarget = new THREE.WebGLRenderTarget( width, height, opt );
  const backgroundTarget = new THREE.WebGLRenderTarget( width, height, opt );

  const materialBg = new THREE.ShaderMaterial( {
    uniforms: {
      u_screen: { value: backgroundTarget.texture },
      u_opacity: { value: WindParam.fadeOpacity },
    },
    vertexShader: quadVert,
    fragmentShader: screenFrag,
    depthTest: false,
    transparent: true,
  } );
  const materialScreen = new THREE.ShaderMaterial( {
    uniforms: {
      u_screen: { value: screenTarget.texture },
      u_opacity: { value: 1.0 },
    },
    vertexShader: quadVert,
    fragmentShader: screenFrag,
    blending: THREE.CustomBlending,
    blendSrc: THREE.OneMinusSrcAlphaFactor, // SrcAlphaFactor OneMinusSrcAlphaFactor
    blendDst: THREE.SrcAlphaFactor,
    blendEquation: THREE.AddEquation,
    depthTest: false,
    transparent: true,
  } );
  // copy from three.js\examples\jsm\postprocessing\Pass.js
  const _camera = new THREE.OrthographicCamera( - 1, 1, 1, - 1, 0, 1 );
  class FullscreenTriangleGeometry extends THREE.BufferGeometry {
    constructor() {
      super();
      this.setAttribute( 'position', new THREE.Float32BufferAttribute( [ - 1, 3, 0, - 1, - 1, 0, 3, - 1, 0 ], 3 ) );
      this.setAttribute( 'uv', new THREE.Float32BufferAttribute( [ 0, 2, 0, 0, 2, 0 ], 2 ) );
    }
  }
  const _geometry = new FullscreenTriangleGeometry()
  // end copy from

  const meshBg = new THREE.Mesh(_geometry, materialBg)
  sceneBack.add( meshBg )
  const meshFront = new THREE.Mesh(_geometry, materialScreen)
  sceneFront.add(meshFront)

  const pointsArr = new Float32Array(_numParticles * 3);
  for (let i = 0; i < pointsArr.length; i++) {
    pointsArr[i] = 0;
  }
  const pointsBuffer = new THREE.Float32BufferAttribute(pointsArr, 3)
  const particleIndices = new Float32Array(_numParticles);
  for (let i = 0; i < _numParticles; i++) {
    particleIndices[i] = i;
  }
  const particleIndexBuffer = new THREE.Float32BufferAttribute(particleIndices, 1)

  const materialDraw = new THREE.ShaderMaterial( {
    uniforms: {
      u_particles: { value: null },
      u_particles_res: { value: particleRes },
      u_wind: { value: windTexture },
      u_color_ramp: { value: textureColors },
      u_wind_min: { value: {x: windData.uMin, y: windData.vMin} },
      u_wind_max: { value: {x: windData.uMax, y: windData.vMax} },
    },
    vertexShader: drawVert,
    fragmentShader: drawFrag,
    depthTest: false,
    transparent: true,
  } );

  const pointsWindGeo = new THREE.BufferGeometry();
  pointsWindGeo.setAttribute('a_index', particleIndexBuffer)
  pointsWindGeo.setAttribute('position', pointsBuffer)
  const pointsWind = new THREE.Points( pointsWindGeo, materialDraw );
  // scene.add( pointsWind );
  sceneBack.add(pointsWind);

  function findRenderTargetByTexture(texture) {
    const targets = [backgroundTarget, screenTarget]
    return targets.find(target => target.texture === texture);
  }
  function findScreenRenderTargetNot(texture) {
    const targets = [backgroundTarget, screenTarget]
    return targets.find(target => target.texture !== texture);
  }

  function animate() {
    positionUniforms[ 'u_rand_seed' ] = { value: Math.random() };
    gpuCompute.compute();

    materialDraw.uniforms.u_particles.value = velocityVariable.material.uniforms.u_particles.value;

    const windTarget = findRenderTargetByTexture(materialScreen.uniforms.u_screen.value);
    renderer.setRenderTarget( windTarget );
    renderer.render( sceneBack, _camera );
    renderer.setRenderTarget( null );

    renderer.render( sceneFront, _camera );
    materialBg.uniforms.u_screen.value = findScreenRenderTargetNot(materialBg.uniforms.u_screen.value).texture;
    materialScreen.uniforms.u_screen.value = findScreenRenderTargetNot(materialScreen.uniforms.u_screen.value).texture;

    planeMaterial.map = materialScreen.uniforms.u_screen.value;

    // planeMaterial.map = velocityVariable.material.uniforms.u_particles.value;

    renderer.render( scene, camera );
    requestAnimationFrame(animate);
  }
  animate()
</script>
</body>
</html>
